//! 内存数据源
//! 
//! 提供基于内存的数据存储和查询功能

use std::collections::HashMap;
use async_trait::async_trait;
use serde_json::{Value, json};

use crate::{DataSource, DataSourceQuery, DataSourceResult, DataSourceStats};

/// 内存数据源
pub struct MemoryDataSource {
    name: String,
    description: String,
    connected: bool,
    data: HashMap<String, Vec<HashMap<String, Value>>>,
}

impl MemoryDataSource {
    /// 创建新的内存数据源
    pub fn new() -> Self {
        let mut data = HashMap::new();
        
        // 添加一些测试数据
        let mut users = Vec::new();
        users.push([
            ("id".to_string(), json!(1)),
            ("name".to_string(), json!("张三")),
            ("age".to_string(), json!(25)),
            ("email".to_string(), json!("zhangsan@example.com")),
        ].iter().cloned().collect());
        
        users.push([
            ("id".to_string(), json!(2)),
            ("name".to_string(), json!("李四")),
            ("age".to_string(), json!(30)),
            ("email".to_string(), json!("lisi@example.com")),
        ].iter().cloned().collect());
        
        users.push([
            ("id".to_string(), json!(3)),
            ("name".to_string(), json!("王五")),
            ("age".to_string(), json!(28)),
            ("email".to_string(), json!("wangwu@example.com")),
        ].iter().cloned().collect());
        
        data.insert("users".to_string(), users);
        
        // 添加产品数据
        let mut products = Vec::new();
        products.push([
            ("id".to_string(), json!(1)),
            ("name".to_string(), json!("笔记本电脑")),
            ("price".to_string(), json!(5999.99)),
            ("category".to_string(), json!("电子产品")),
        ].iter().cloned().collect());
        
        products.push([
            ("id".to_string(), json!(2)),
            ("name".to_string(), json!("智能手机")),
            ("price".to_string(), json!(2999.99)),
            ("category".to_string(), json!("电子产品")),
        ].iter().cloned().collect());
        
        data.insert("products".to_string(), products);
        
        Self {
            name: "内存数据源".to_string(),
            description: "基于内存的数据存储，用于测试和演示".to_string(),
            connected: false,
            data,
        }
    }
    
    /// 添加表数据
    pub fn add_table(&mut self, table_name: String, rows: Vec<HashMap<String, Value>>) {
        self.data.insert(table_name, rows);
    }
    
    /// 获取表数据
    pub fn get_table(&self, table_name: &str) -> Option<&Vec<HashMap<String, Value>>> {
        self.data.get(table_name)
    }
    
    /// 列出所有表
    pub fn list_tables(&self) -> Vec<&str> {
        self.data.keys().map(|k| k.as_str()).collect()
    }
    
    /// 简单的查询解析和执行
    fn execute_simple_query(&self, query: &str) -> anyhow::Result<Value> {
        let query = query.trim().to_lowercase();
        
        if query.starts_with("select * from ") {
            let table_name = query.strip_prefix("select * from ").unwrap().trim();
            if let Some(table_data) = self.data.get(table_name) {
                Ok(json!({
                    "rows": table_data,
                    "count": table_data.len()
                }))
            } else {
                Err(anyhow::anyhow!("表不存在: {}", table_name))
            }
        } else if query.starts_with("show tables") {
            let tables: Vec<_> = self.data.keys().collect();
            Ok(json!({
                "tables": tables
            }))
        } else if query.starts_with("describe ") {
            let table_name = query.strip_prefix("describe ").unwrap().trim();
            if let Some(table_data) = self.data.get(table_name) {
                if let Some(first_row) = table_data.first() {
                    let columns: Vec<_> = first_row.keys().collect();
                    Ok(json!({
                        "table": table_name,
                        "columns": columns
                    }))
                } else {
                    Ok(json!({
                        "table": table_name,
                        "columns": []
                    }))
                }
            } else {
                Err(anyhow::anyhow!("表不存在: {}", table_name))
            }
        } else {
            Err(anyhow::anyhow!("不支持的查询: {}", query))
        }
    }
}

impl Default for MemoryDataSource {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl DataSource for MemoryDataSource {
    fn name(&self) -> &str {
        &self.name
    }
    
    fn description(&self) -> &str {
        &self.description
    }
    
    async fn connect(&mut self) -> anyhow::Result<()> {
        tracing::info!("连接到内存数据源");
        self.connected = true;
        Ok(())
    }
    
    async fn disconnect(&mut self) -> anyhow::Result<()> {
        tracing::info!("断开内存数据源连接");
        self.connected = false;
        Ok(())
    }
    
    async fn is_connected(&self) -> bool {
        self.connected
    }
    
    async fn execute_query(&mut self, query: DataSourceQuery) -> anyhow::Result<DataSourceResult> {
        if !self.connected {
            return Err(anyhow::anyhow!("数据源未连接"));
        }
        
        let start_time = std::time::Instant::now();
        
        tracing::debug!("执行内存查询: {}", query.content);
        
        match self.execute_simple_query(&query.content) {
            Ok(data) => {
                let duration = start_time.elapsed();
                let rows_processed = if let Some(rows) = data.get("rows") {
                    rows.as_array().map(|arr| arr.len() as u64).unwrap_or(0)
                } else {
                    1
                };
                
                Ok(DataSourceResult {
                    success: true,
                    data,
                    error: None,
                    stats: DataSourceStats {
                        execution_time_ms: duration.as_millis() as f64,
                        rows_processed,
                        bytes_processed: query.content.len() as u64,
                    },
                })
            }
            Err(e) => {
                let duration = start_time.elapsed();
                Ok(DataSourceResult {
                    success: false,
                    data: json!({}),
                    error: Some(e.to_string()),
                    stats: DataSourceStats {
                        execution_time_ms: duration.as_millis() as f64,
                        rows_processed: 0,
                        bytes_processed: query.content.len() as u64,
                    },
                })
            }
        }
    }
    
    async fn health_check(&self) -> anyhow::Result<Value> {
        Ok(json!({
            "status": if self.connected { "connected" } else { "disconnected" },
            "name": self.name,
            "tables": self.list_tables(),
            "total_rows": self.data.values().map(|v| v.len()).sum::<usize>()
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_memory_data_source() {
        let mut source = MemoryDataSource::new();
        
        // 测试连接
        assert!(!source.is_connected().await);
        source.connect().await.unwrap();
        assert!(source.is_connected().await);
        
        // 测试查询
        let query = DataSourceQuery {
            id: "test".to_string(),
            content: "SELECT * FROM users".to_string(),
            parameters: HashMap::new(),
        };
        
        let result = source.execute_query(query).await.unwrap();
        assert!(result.success);
        assert!(result.stats.rows_processed > 0);
        
        // 测试健康检查
        let health = source.health_check().await.unwrap();
        assert_eq!(health["status"], "connected");
        
        // 测试断开连接
        source.disconnect().await.unwrap();
        assert!(!source.is_connected().await);
    }
    
    #[tokio::test]
    async fn test_memory_queries() {
        let mut source = MemoryDataSource::new();
        source.connect().await.unwrap();
        
        // 测试SHOW TABLES
        let query = DataSourceQuery {
            id: "test".to_string(),
            content: "SHOW TABLES".to_string(),
            parameters: HashMap::new(),
        };
        
        let result = source.execute_query(query).await.unwrap();
        assert!(result.success);
        
        // 测试DESCRIBE
        let query = DataSourceQuery {
            id: "test".to_string(),
            content: "DESCRIBE users".to_string(),
            parameters: HashMap::new(),
        };
        
        let result = source.execute_query(query).await.unwrap();
        assert!(result.success);
    }
}
